import base64
import re
from agent.agent_selector import determine_agent
from agent.content_creator_agent import content_creator_tool
from agent.content_draw_agent import content_draw_tool
from agent.travel_draw_agent import travel_draw_tool, generate_image
from agent.travel_assistant_agent import travel_assistant_tool
from tts.text_to_speech import generate_audio_from_text

image_path = 'src/assets/flow.png'


def convert_image_to_base64(image_path):
    # 读取图片文件内容
    with open(image_path, 'rb') as image_file:
        image_data = image_file.read()
    # 转换为 Base64 编码
    image_base64 = base64.b64encode(image_data).decode('utf-8')
    return image_base64


def get_agent_response(user_input):
    img_base64 = None
    agent_type, model_response = determine_agent(user_input)
    day_images = []  # 初始化为空列表，以防变量未定义的情况
    mp3_base64 = None  # 初始化为空，防止变量未定义

    if agent_type == "travel_assistant":
        response_text = travel_assistant_tool.func(user_input)
        role_used = "ERNIE-Character-8K+Spark Lite+Stable-Diffusion-XL"

        # 移除或替换所有星号(*)，以避免潜在的问题
        response_text = response_text.replace('*', '')

        mp3_base64 = generate_audio_from_text(response_text)
        img_base64 = convert_image_to_base64(image_path)
        # 使用正则表达式按关键词分段，并保留分段关键词
        #print("推理base64"+img_base64)
        itinerary_sections = re.split(r'(\bDay \d+|First|Next|Then|Finally|Lastly)\b', response_text)
        itinerary_sections = [section.strip() for section in itinerary_sections if section.strip()]

        # 检查是否找到分段关键词
        if len(itinerary_sections) < 1:
            itinerary_sections = [section.strip() for section in response_text.split("\n\n") if section.strip()]

        # 组合关键词与内容
        combined_sections = []
        current_text = ""
        for i, section in enumerate(itinerary_sections):
            if re.match(r'^(Day \d+|First|Next|Then|Finally|Lastly)$', section):
                if current_text:
                    combined_sections.append(current_text)
                current_text = section  # 关键词作为新段落的开头
            else:
                current_text += " " + section  # 把内容拼接到关键词之后
        if current_text:
            combined_sections.append(current_text)  # 追加最后一段

        day_started = False

        for section in combined_sections:
            # 检查是否遇到第一个 Day 段落
            if not day_started and re.match(r'^(Day \d+|First)', section, re.IGNORECASE):
                day_started = True  # 第一个 Day 段落开始后，标记为 True

            # 生成绘图提示词和图像
            if day_started:
                drawing_prompt = travel_draw_tool.func(section)
                print("绘图提示词:", drawing_prompt)

                image_base64 = generate_image(drawing_prompt)
                print("图像编码:", image_base64)

                # 添加段落到 day_images，段落后跟一个空行
                day_images.append({
                    "text": f"{section}\n",
                    "image": image_base64
                })
            else:
                # 在 Day1 之前的段落仅添加文本，不生成图像
                day_images.append({
                    "text": f"{section}\n",
                    "image": None
                })

    elif agent_type == "content_creator":
        response_text = content_creator_tool.func(user_input)
        role_used = "ERNIE-8K+Spark Lite+Stable-Diffusion-XL"
        response_text = response_text.replace('*', '')
        mp3_base64 = generate_audio_from_text(response_text)
        img_base64 = convert_image_to_base64(image_path)
        content_sections = [section.strip() for section in response_text.split("\n\n") if section.strip()]
        print("推理base64"+img_base64)
        # 根据段落数量决定生成图像的策略
        num_sections = len(content_sections)
        if num_sections <= 4:
            sections_to_draw = [content_sections[i:i + 2] for i in range(0, num_sections, 2)]
        else:
            sections_to_draw = [content_sections[i:i + 5] for i in range(0, num_sections, 5)]

        for sections in sections_to_draw:
            # 合并段落文本作为绘图提示词
            combined_section = "\n\n".join(sections)
            drawing_prompt = content_draw_tool.func(combined_section)
            print("绘图提示词:", drawing_prompt)

            # 使用生成的提示词生成图像
            image_base64 = generate_image(drawing_prompt)
            print("图像编码:", image_base64)

            # 将合并后的文本和图像组合后添加到 `day_images`
            day_images.append({
                "text": f"{combined_section}\n",  # 不再加粗首单词，并添加空行
                "image": image_base64
            })


    else:
        response_text = model_response
        role_used = "无法确定角色"

    return {
        "role_used": role_used,
        "days": day_images,  # 每天的文本和图像组合
        "result_audio_base64": mp3_base64,  # 整个回复的音频Base64编码
        "img_base64": img_base64
    }
